from enthought.traits.api import HasTraits,Int,Float,Str,Property
from enthought.traits.api import Range,Array, Instance
from enthought.traits.ui.api import View,Item,Label
from chaco.api import Plot, ArrayPlotData
from chaco import default_colors as dc
from enable.component_editor import ComponentEditor
import numpy as np
import filereader as fr

class WaterfallPlotter(HasTraits):
    '''Creates a waterfall plot from 2D datasets.'''
    
    data = {}

    # Creates some sample data for testing.
    xslider = Range(0.0, 20.0)
    yslider = Range(800.0, 5000.0)

    xoffset = Property(Float, depends_on=["xslider"])
    yoffset = Property(Float, depends_on=["yslider"])
    
    plot = Instance(Plot)

    # Sets default values.
    def _xslider_default(self): return 1.0
    def _yslider_default(self): return 1000.0

    def _get_xoffset(self): return self.xslider 
    def _get_yoffset(self): return self.yslider

    traits_view = View(
             Item('plot', editor=ComponentEditor(), show_label=False, ),
             Item('xslider', label='X Offset'),
             Item('yslider', label='Y Offset'), 
             width=900, height=600, resizable=True, title="Waterfall Plotter")

    def __init__(self, data):
        super(WaterfallPlotter, self).__init__()
        self.data = data
        
        self.graph()

        self.on_trait_change(self.graph, 'xslider')
        self.on_trait_change(self.graph, 'yslider')

    def graph(self):
        plotdata = ArrayPlotData()
        plot = Plot(plotdata)

        cp = 0
        cr = 0
        # Load datasets into the ArrayPlotData object.
        for i, key in enumerate(self.data):
            xlabel = "x%d" % key
            ylabel = "y%d" % key

            xdata = self.data[key][0]+i*self.xoffset
            ydata = self.data[key][1]+i*self.yoffset

            plotdata.set_data(xlabel, xdata)
            plotdata.set_data(ylabel, ydata)
            plot.plot((xlabel, ylabel), type="line", 
                      color=tuple(dc.PALETTES[cp%3][cr%10]))
            cp += 1
            cr += 1
 
        plot.title = "Title Here"
        self.plot = plot

if __name__ == "__main__":
    data = fr.getDataSets('.chi')
    p = WaterfallPlotter(data)
    p.configure_traits()

